#include "robot_sim/robot_webots_driver.hpp"

#include <iostream>
#include <memory>
#include <webots/gyro.h>
#include <webots/inertial_unit.h>
#include <webots/motor.h>
#include <webots/position_sensor.h>
#include <webots/robot.h>

namespace webots_driver
{
void WebotsDriver::init(
    webots_ros2_driver::WebotsNode *node,
    std::unordered_map<std::string, std::string>
        &parameters)
{
    node->declare_parameter<bool>("direct_control", false);
    node->declare_parameter<double>("forward_scale", 1.0);
    node->declare_parameter<double>("rotation_scale", 1.0);
    directControl =
        node->get_parameter("direct_control").as_bool();
    forwardScale =
        node->get_parameter("forward_scale").as_double();
    rotationScale =
        node->get_parameter("rotation_scale").as_double();
    std::string robotName = wb_robot_get_name();
    if (directControl)
    {
        RCLCPP_INFO(node->get_logger(),
                    "WebotsDriver: direct control mode");
        sub = node->create_subscription<ControlSignal>(
            robotName + "/control_signal", 1,
            [this](const ControlSignal &signal)
            {
                forwardInput =
                    signal.forward_control * forwardScale;
                rotationInput =
                    signal.rotate_control * rotationScale;
            });
    }
    else
    {
        RCLCPP_INFO(node->get_logger(),
                    "WebotsDriver: servo mode");
        servoSub = node->create_subscription<ChassisServo>(
            robotName + "/chassis_servo", 1,
            std::bind(&WebotsDriver::OnChassisServo, this,
                      std::placeholders::_1));
    }
    rightWheel = wb_robot_get_device("right_wheel_motor");
    leftWheel = wb_robot_get_device("left_wheel_motor");
    rightSensor = wb_motor_get_position_sensor(rightWheel);
    leftSensor = wb_motor_get_position_sensor(leftWheel);
    wb_position_sensor_enable(leftSensor, 5);
    wb_position_sensor_enable(rightSensor, 5);
    lastLeftPosition =
        wb_position_sensor_get_value(leftSensor) * 180 /
        M_PI;
    lastRightPosition =
        wb_position_sensor_get_value(rightSensor) * 180 /
        M_PI;
    leftSensor = wb_motor_get_position_sensor(rightWheel);
    imu = wb_robot_get_device("imu");
    gyro = wb_robot_get_device("gyro");
    wb_motor_set_position(rightWheel, INFINITY);
    wb_motor_set_position(leftWheel, INFINITY);
    wb_motor_set_velocity(rightWheel, 0);
    wb_motor_set_velocity(leftWheel, 0);

    wb_gyro_enable(gyro, 5);
    wb_inertial_unit_enable(imu, 5);
    if (parameters.contains("gyro_PID"))
    {
        gyro_pid = createPID(parameters["gyro_PID"]);
    }
    if (parameters.contains("angle_PID"))
    {
        angle_pid = createPID(parameters["angle_PID"]);
    }
    if (parameters.contains("speed_PID"))
    {
        speed_pid = createPID(parameters["speed_PID"]);
    }
    if (parameters.contains("rotation_PID"))
    {
        rotation_pid =
            createPID(parameters["rotation_PID"]);
    }
}

std::unique_ptr<CtrlAlg::PIDPos<double>> WebotsDriver::
    createPID(const std::string &param)
{
    std::stringstream ss(param);
    double params[3];
    for (int i = 0; i < 3; i++)
    {
        ss >> params[i];
    }
    return std::make_unique<CtrlAlg::PIDPos<double>>(
        params[0], params[1], params[2]);
}

void WebotsDriver::step()
{
    if (directControl)
    {
        DirectControl();
    }
    else
    {
        ChassisServoControl();
    }
}

void WebotsDriver::DirectControl()
{
    const double *angles =
        wb_inertial_unit_get_roll_pitch_yaw(imu);
    double angle = -angles[1] * 180 / M_PI;
    // double angleZ = angles[2] * 180 / M_PI;
    const double *gyroes = wb_gyro_get_values(gyro);
    double gyro = -gyroes[1] * 180 / M_PI;
    double gyroZ = gyroes[2] * 180 / M_PI;
    double poseLeft =
        wb_position_sensor_get_value(leftSensor) * 180 /
        M_PI;
    double wheelSpeedLeft = poseLeft - lastLeftPosition;
    lastLeftPosition = poseLeft;
    double poseRight =
        wb_position_sensor_get_value(rightSensor) * 180 /
        M_PI;
    double wheelSpeedRight = poseRight - lastRightPosition;
    lastRightPosition = poseRight;

    double speed = (wheelSpeedLeft + wheelSpeedRight) / 2;
    double angleCtrl =
        speed_pid->GetControl(forwardInput - speed, 0.01);
    angleCtrl = std::clamp<double>(angleCtrl, -45, 45);
    // for test angle loop and gyro loop
    // angleCtrl = forwardInput;
    double gyroCtrl =
        angle_pid->GetControl(angleCtrl - angle, 0.01);
    gyroCtrl = std::clamp<double>(gyroCtrl, -10, 10);
    // for test gyro loop only
    // gyroCtrl = forwardInput;
    double lastControl =
        -gyro_pid->GetControl(gyroCtrl - gyro, 0.01);
    double rotateCtrl = -rotation_pid->GetControl(
        rotationInput - gyroZ, 0.01);
    lastControl = std::clamp<double>(lastControl, -50, 50);
    double lastControlLeft = lastControl - rotateCtrl;
    double lastControlRight = lastControl + rotateCtrl;
    wb_motor_set_torque(leftWheel, lastControlLeft);
    wb_motor_set_torque(rightWheel, lastControlRight);

    // gyro loop only
    // std::cout << "target" << forwardInput << "\tgyro" <<
    // gyro << "\tcontrol"
    // << lastControl << std::endl; angle loop and gyro loop
    // std::cout << "target" << forwardInput << "\tgc" <<
    // gyroCtrl << "\tgyro" << gyro <<
    // "\tangle" << angle <<
    // "\tcontrol" << lastControl << std::endl;
    // std::cout << "target" << forwardInput << "\tspeed" <<
    // speed << "\tangle"
    // << angle << "\tcontrol" << lastControl
    //           << std::endl;
}

void WebotsDriver::ChassisServoControl()
{
    const double *angles =
        wb_inertial_unit_get_roll_pitch_yaw(imu);
    double angle = -angles[0] * 180 / M_PI;
    // double angleZ = angles[2] * 180 / M_PI;
    const double *gyroes = wb_gyro_get_values(gyro);
    double gyro = -gyroes[0] * 180 / M_PI;
    double gyroZ = gyroes[2] * 180 / M_PI;
    double poseLeft =
        wb_position_sensor_get_value(leftSensor) * 180 /
        M_PI;
    double wheelSpeedLeft = poseLeft - lastLeftPosition;
    lastLeftPosition = poseLeft;
    double poseRight =
        wb_position_sensor_get_value(rightSensor) * 180 /
        M_PI;
    double wheelSpeedRight = poseRight - lastRightPosition;
    lastRightPosition = poseRight;

    double angleCtrl = 0;
    if (target.speedGain > 0)
    {
        double speed =
            (wheelSpeedLeft + wheelSpeedRight) / 2;
        angleCtrl = speed_pid->GetControl(
            target.speed - speed, 0.01);
    }
    else if (target.angleGain > 0)
    {
        angleCtrl = target.angle;
    }
    angleCtrl = std::clamp<double>(angleCtrl, -45, 45);
    double gyroCtrl =
        angle_pid->GetControl(angleCtrl - angle, 0.01);
    gyroCtrl = std::clamp<double>(gyroCtrl, -10, 10);
    double lastControl =
        -gyro_pid->GetControl(gyroCtrl - gyro, 0.01);
    double rotateCtrl = -rotation_pid->GetControl(
        target.gyroZ - gyroZ, 0.01);
    lastControl = std::clamp<double>(lastControl, -50, 50);
    double lastControlLeft = lastControl - rotateCtrl;
    double lastControlRight = lastControl + rotateCtrl;
    wb_motor_set_torque(leftWheel, lastControlLeft);
    wb_motor_set_torque(rightWheel, lastControlRight);
}

void WebotsDriver::OnChassisServo(const ChassisServo &msg)
{
    target.angle = msg.angle;
    target.gyro = msg.gyro;
    target.speed = msg.speed;
    target.gyroZ = msg.gyroz;

    target.angleGain = msg.angle_gain;
    target.gyroGain = msg.gyro_gain;
    target.speedGain = msg.speed_gain;
    target.gyroZGain = msg.gyroz_gain;
}

} // namespace webots_driver
#include "pluginlib/class_list_macros.hpp"
PLUGINLIB_EXPORT_CLASS(webots_driver::WebotsDriver,
                       webots_ros2_driver::PluginInterface)